feat: add graceful shutdown with signal handling #174

ben11211 · 2025-09-11T01:29:01Z

Allows for graceful shutdown of the IsolateServicer and associated agents.

Adds the ISOLATE_SHUTDOWN_GRACE_PERIOD environment variable. Defaults to 60m.

Prior Behavior

IsolateServicers client grpc connection closes, but this does not force an exit.
termination interceptor is called which stops the server, but does not terminate agent tasks. Agent tasks exit when they exit, and then main returns.

Proposed Behavior

IsolateServicer is initialized with signal handlers
Shutdown event triggered via either of the following
- Signal handlers receive SIGTERM or SIGINT
- IsolateServicer's gRPC client is no longer active during single-run mode
servicer.initiate_shutdown() called
Servicer calls agent.Terminate() for all servicer's active background_tasks
agent.Terminate SIGTERMs agent process pid by terminating self._bound_context, and SIGKILL after configured grace period

View client termination behavior before and after change.

shutdownbehavior.mov

⚠️ Note ⚠️

Grace period handling is not complete, because agent signal handling can not be propagated completely. When the agent subprocess receives a termination signal, the signal can only be received by the main execution thread of the python application which in this case is src/isolate/connections/grpc/agent.py:run_agent which is a running gRPC server. Agent gRPC calls are handled by worker threads, and the Run handler calls the provided function within the context of the handler thread. This means any provided functions being executed within an isolate agent are not run by a thread that is capable of receiving signals.

This means, within the current multithreaded agent approach - in order to provide graceful shutdown capabilities to isolated python functions, they must have an application layer signal such as a passed context. This requires an function running within an Isolate agent to be context aware if it wants to implement graceful shutdown.

Alternatively modify the agent gRPC server to use AsyncIO instead of the current multithreaded approach. This would not have a performance hit, since the current approach is configured for a single worker thread. Isolate functions can register signal handlers as they would normally without needing to be context aware.

Here's a rough sketch of what that would look like, with a test confirming the signal propagation through to execute_function does work.

src/isolate/connections/grpc/_base.py

src/isolate/server/server.py

ben11211 · 2025-09-12T22:27:45Z

src/isolate/server/server.py


-        server.add_insecure_port("[::]:50001")
-        print("Started listening at localhost:50001")
+        server.add_insecure_port(f"[::]:{options.port}")


This change was necessary to support end to end testing concurrently.

tests/test_connections.py

tests/test_shutdown.py

…isconnect detection

tests/test_shutdown.py

chamini2 · 2025-09-12T23:22:15Z

src/isolate/connections/grpc/_base.py

+        try:
+            print(f"Terminating agent PID {proc.pid}")
+            proc.terminate()
+            proc.wait(timeout=shutdown_grace_period)
+        except subprocess.TimeoutExpired:
+            # Process didn't die within timeout
+            print(f"killing agent PID {proc.pid}")
+            proc.kill()


looking good

chamini2 · 2025-09-12T23:31:25Z

src/isolate/server/server.py

+    def initiate_shutdown(self, grace_period: float | None = None) -> None:
+        if self._shutting_down:
+            return
+        self._shutting_down = True
+        if grace_period is None:
+            grace_period = SHUTDOWN_GRACE_PERIOD


Accepting grace_period as None reads as if it was the "infinite wait" option, but in reality you are assigning a default of SHUTDOWN_GRACE_PERIOD.

I would change grace_period to directly default to the SHUTDOWN_GRACE_PERIOD in the function params

Instead, I've removed it from the function parameters. I realizes it was leftover from tests in a previous revision, when the terminate -> kill timing was managed within initiate_shutdown. Now this function only outputs a useful log with this value.

SHUTDOWN_GRACE_PERIOD is now passed through to the LocalPythonGRPC class during agent allocation, and it controls the terminate -> kill timing there instead.

chamini2 · 2025-09-12T23:33:40Z

src/isolate/server/server.py

+        # Collect all active agents from running tasks
+        shutdown_threads = []
+        for task in self.background_tasks.values():
+            if task.agent is not None:


and if it is None you should probably cancel it directly

chamini2 · 2025-09-12T23:36:16Z

src/isolate/server/server.py

+
+        if self._server:
+            print("Stopping gRPC server")
+            self._server.stop(grace=0.1)  # Short grace period for server shutdown


I think bringing the server into the servicer context has no real benefit, you can call initiate_shutdown and call server.stop() right after in the SingleTaskInterceptor context

chamini2 · 2025-09-12T23:41:29Z

src/isolate/server/server.py

+    def register_signal_handlers(self, server: grpc.Server) -> None:
+        """Set up signal handlers for graceful shutdown"""
+        self._server = server
+
+        def signal_handler(signum, _):
+            """Handle SIGTERM and SIGINT by gracefully shutting down server"""
+            print(f"Received signal {signum}, shutting down server")
+            self.initiate_shutdown()
+
+        signal.signal(signal.SIGTERM, signal_handler)
+        signal.signal(signal.SIGINT, signal_handler)


And then you would need to move this function out of this context too, just to a context that has access to both the server and the servicer

you can just do right before the main function a

def signal_shutdown(signum, _): print(f"Received signal {signum}, shutting down server") servicer.initiate_shutdown() server.stop() def main(): # ... signal.signal(signal.SIGTERM, signal_shutdown) signal.signal(signal.SIGINT, signal_shutdown)

chamini2 · 2025-09-12T23:49:41Z

tests/test_connections.py

+        # Process should be terminated after terminate_proc() returns
+        assert proc.poll() is not None, "Process should be terminated by SIGTERM"
+
+    def test_force_terminate(self):


chamini2 · 2025-09-12T23:52:16Z

tests/test_shutdown.py

+    def function_obj():
+        import time
+
+        time.sleep(10)


should this be an infinite loop instead? to make sure it is just not taking 10 secs to finish the test but it is actually closing because of the disconnect

…isconnect detection

ben11211 · 2025-09-13T04:23:09Z

tests/test_shutdown.py

+def create_run_request(func, stream_logs=True):
+    """Convert a Python function into a BoundFunction request for stub.Run()."""
+    bound_function = functools.partial(func)
+    serialized_function = to_serialized_object(bound_function, method="cloudpickle")
+
+    env_def = EnvironmentDefinition()
+    env_def.kind = "local"
+
+    request = BoundFunction()
+    request.function.CopyFrom(serialized_function)
+    request.environments.append(env_def)
+    request.stream_logs = stream_logs
+
+    return request


@cmlad Here's a convenience function for passing normal functions to an IsolateServicer.Run call. FYI.

ben11211 · 2025-09-13T04:28:52Z

src/isolate/server/server.py

+def register_signal_handlers(servicer: IsolateServicer, server: grpc.Server) -> None:
+    def handle_signal(signum, frame):
+        print(f"Received signal {signum}, initiating shutdown...")
+        servicer.initiate_shutdown()
+        server.stop(grace=0.1)
+
+    signal.signal(signal.SIGINT, handle_signal)
+    signal.signal(signal.SIGTERM, handle_signal)


@chamini2 Is this what you're suggesting here #174 (comment)?

ben11211 changed the title ~~feat(server): add graceful shutdown with signal handling and client d…~~ feat: add graceful shutdown with signal handling Sep 11, 2025

ben11211 force-pushed the server-agent-graceful-shutdown branch 5 times, most recently from c3cb5ef to b86f35d Compare September 12, 2025 07:45

ben11211 marked this pull request as ready for review September 12, 2025 07:47

ben11211 force-pushed the server-agent-graceful-shutdown branch from b86f35d to 3ac8323 Compare September 12, 2025 20:40